import mxnet as mx
from mxnet.gluon import nn
import gluoncv,gluonnlp
import mxnet.ndarray as nd
class Resnet_CTC(mx.gluon.Block):
    def __init__(self,sequence_len = 152, alphabet_size = 300):
        super(Resnet_CTC,self).__init__()
        self.features = gluoncv.model_zoo.resnet18_v1b(pretrained=True,dilated = True)
        self.ctc = mx.gluon.loss.CTCLoss(layout='NTC', label_layout='NT')
        self.conv = nn.Conv2D(channels=sequence_len * alphabet_size , kernel_size=(1, 1), strides=(1, 1))
        self.conv.bias.initialize(init=mx.init.Zero())
        self.conv.weight.initialize(init=mx.init.Normal())

        self._sequence_len = sequence_len
        self._alphabet_size = alphabet_size
    def forward(self, x):
        feat = self.features
        x = feat.conv1(x)
        x = feat.bn1(x)
        x = feat.relu(x)
        x = feat.maxpool(x)

        x = feat.layer1(x)
        x = feat.layer2(x)
        x = feat.layer3(x)
        x = feat.layer4(x)
        x = self.conv(x)
        x = nd.reshape(x,(x.shape[0],x.shape[1],-1)) # (n, sequence_len, w, h)
        x = nd.max(x,axis=2)
        x = nd.reshape(x,(x.shape[0],self._sequence_len,self._alphabet_size))
        return x

if __name__ == '__main__':
    res = Resnet_CTC()
    res.collect_params().reset_ctx(mx.gpu(2))
    x= mx.nd.zeros(shape=(1,3,64,256),ctx=mx.gpu(2))
    y = res(x)
    print(y.shape)
